package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"strconv"
	"strings"

	"github.com/segmentio/kafka-go"
)

var (
	// 命令行参数
	// 新建topic
	createTopicCmd = &flagSet{
		FlagSet:    flag.NewFlagSet("create", flag.ExitOnError),
		cmdComment: "create topic",
	}
	createbrokers = createTopicCmd.String("brokers", "localhost:9092", "kafka brokers")
	createtopic   = createTopicCmd.String("topic", "test", "kafka topic")
	numPartitions = createTopicCmd.Int("partitions", 1, "kafka partitions")
	replication   = createTopicCmd.Int("replication", 1, "kafka replication factor")

	// 列出topic
	listTopicCmd = &flagSet{
		FlagSet:    flag.NewFlagSet("list", flag.ExitOnError),
		cmdComment: "list topics",
	}
	listbrokers = listTopicCmd.String("brokers", "localhost:9092", "kafka brokers")

	// 对topic消费
	consumerCmd = &flagSet{
		FlagSet:    flag.NewFlagSet("consumer", flag.ExitOnError),
		cmdComment: "consumer messages",
	}
	consumerbrokers = consumerCmd.String("brokers", "localhost:9092", "kafka brokers")
	consumertopic   = consumerCmd.String("topic", "test", "kafka topic")
	consumerGroup   = consumerCmd.String("group", "test", "kafka consumer group")

	// 删除topic
	deleteTopicCmd = &flagSet{
		FlagSet:    flag.NewFlagSet("delete", flag.ExitOnError),
		cmdComment: "delete topic",
	}
	deletebrokers = deleteTopicCmd.String("brokers", "localhost:9092", "kafka brokers")
	deletetopic   = deleteTopicCmd.String("topic", "test", "kafka topic")

	// 子命令
	subcommands = map[string]*flagSet{
		createTopicCmd.Name(): createTopicCmd,
		consumerCmd.Name():    consumerCmd,
		listTopicCmd.Name():   listTopicCmd,
		deleteTopicCmd.Name(): deleteTopicCmd,
	}
	cmd = subcommands[os.Args[1]]
)

type flagSet struct {
	*flag.FlagSet
	cmdComment string // 二级子命令本身的注释
}

func main() {
	useage := func() { // 整个命令行的帮助信息
		fmt.Printf("Usage: kafka-tool COMMAND\n\n")
		for _, v := range subcommands {
			fmt.Printf("%s: %s\n", v.Name(), v.cmdComment)
			v.PrintDefaults() // 使用 flag 库自带的格式输出子命令的选项帮助信息
			fmt.Println()
		}
		os.Exit(2)
	}

	if len(os.Args) < 2 { // 即没有输入子命令
		useage()
	}

	// 第二个参数必须是我们支持的子命令
	if cmd == nil {
		useage()
	}

	cmd.Parse(os.Args[2:])
	flag.Parse()
	fmt.Println("command name is:", cmd.Name())
	cmd.Visit(func(f *flag.Flag) {
		fmt.Printf("option %s: %s\n", f.Name, f.Value)
	})
	switch cmd.Name() {
	case consumerCmd.Name():
		reader()
	case listTopicCmd.Name():
		listTopic()
	case createTopicCmd.Name():
		createTopic()
	case deleteTopicCmd.Name():
		deleteTopic()
	}
}

func deleteTopic() {
	for _, broker := range strings.Split(*deletebrokers, ",") {
		conn, err := kafka.Dial("tcp", broker)
		if err != nil {
			//panic(err.Error())
			continue
		}
		defer conn.Close()

		controller, err := conn.Controller()
		if err != nil {
			//panic(err.Error())
			continue
		}
		var controllerConn *kafka.Conn
		controllerConn, err = kafka.Dial("tcp", net.JoinHostPort(controller.Host, strconv.Itoa(controller.Port)))
		if err != nil {
			//panic(err.Error())
			continue
		}
		defer controllerConn.Close()

		err = controllerConn.DeleteTopics(*deletetopic)
		if err != nil {
			panic(err.Error())
		}
		break
	}
}

func createTopic() {
	// to create topics when auto.create.topics.enable='false'
	for _, broker := range strings.Split(*createbrokers, ",") {
		conn, err := kafka.Dial("tcp", broker)
		if err != nil {
			//panic(err.Error())
			continue
		}
		defer conn.Close()

		controller, err := conn.Controller()
		if err != nil {
			//panic(err.Error())
			continue
		}
		var controllerConn *kafka.Conn
		controllerConn, err = kafka.Dial("tcp", net.JoinHostPort(controller.Host, strconv.Itoa(controller.Port)))
		if err != nil {
			//panic(err.Error())
			continue
		}
		defer controllerConn.Close()

		topicConfigs := []kafka.TopicConfig{
			{
				Topic:             *createtopic,
				NumPartitions:     *numPartitions,
				ReplicationFactor: *replication,
			},
		}

		err = controllerConn.CreateTopics(topicConfigs...)
		if err != nil {
			panic(err.Error())
		}
	}

}

func listTopic() {
	for _, broker := range strings.Split(*listbrokers, ",") {
		conn, err := kafka.Dial("tcp", broker)
		if err != nil {
			continue
			//panic(err.Error())
		}
		defer conn.Close()

		partitions, err := conn.ReadPartitions()
		if err != nil {
			//panic(err.Error())
			continue
		}

		m := map[string]struct{}{}

		for _, p := range partitions {
			m[p.Topic] = struct{}{}
		}
		for k := range m {
			fmt.Println(k)
		}
		break
	}

}

func reader() {
	r := kafka.NewReader(kafka.ReaderConfig{
		Brokers:  strings.Split(*consumerbrokers, ","),
		GroupID:  *consumerGroup,
		Topic:    *consumertopic,
		MaxBytes: 10e6, // 10MB
	})

	for {
		m, err := r.ReadMessage(context.Background())
		if err != nil {
			fmt.Printf("failed to read message: %v\n", err)
			break
		}
		fmt.Printf("message at offset %d: %s = %s\n", m.Offset, string(m.Key), string(m.Value))
	}

	if err := r.Close(); err != nil {
		log.Fatal("failed to close reader:", err)
	}
}
